import numpy as np
import torch
from scipy.spatial.distance import directed_hausdorff


def dice_coefficient_3D(predict, target, num_classes, epsilon=1e-6):
    # The first output of this network has been reshaped to (-1,class_num). Softmax operation is also performed

    if isinstance(predict, torch.Tensor) and isinstance(target, torch.Tensor):
        total_dice = 0
        for i in range(num_classes):
            # if len(target.shape) > 2:
            #     target = target[:, i, :, :, :].reshape(-1, 1)
            # if len(predict.shape) > 2:
            #     predict = predict[:, i, :, :, :].reshape(-1, 1)
            intersection = torch.sum(torch.mul(predict, target))
            debug_predict_sum = torch.sum(predict)
            debug_target_sum = torch.sum(target)
            dice = (2 * intersection + epsilon) / (debug_predict_sum + debug_target_sum + epsilon)
            total_dice += dice
        return total_dice / num_classes
    elif isinstance(predict, np.ndarray) and isinstance(target, np.ndarray):
        total_dice = 0
        for i in num_classes:
            target = target[:, i, :]
            predict = predict[:, i, :]
            intersection = np.sum(np.multiply(predict, target))
            debug_predict_sum = np.sum(predict)
            debug_target_sum = np.sum(target)
            dice = (2 * intersection + epsilon) / (debug_predict_sum + debug_target_sum + epsilon)
            total_dice += dice
        return total_dice / num_classes
    else:
        raise TypeError


def Hausdorff_Distance_3D(predict, target):
    assert (predict.shape == target.shape)
    if isinstance(predict, torch.Tensor):
        predict = predict.clone().detach().cpu().numpy()
    if isinstance(target, torch.Tensor):
        target = target.clone().detach().cpu().numpy()
    Hausdorff_Distance: float = 0.
    compute_num = 0
    if len(predict.shape) == 5:
        for batch_num in range(predict.shape[0]):
            for channel_num in range(predict.shape[1]):
                for z_index in range(predict.shape[4]):
                    one_layer_HD = directed_hausdorff(predict[batch_num, channel_num, :, :, z_index],
                                                      target[batch_num, channel_num, :, :, z_index])[0]
                    compute_num += 1
                    Hausdorff_Distance += one_layer_HD
        return Hausdorff_Distance, round(Hausdorff_Distance / compute_num, 2)

    else:
        for z_index in range(predict.shape[2]):
            one_layer_HD = directed_hausdorff(predict[:, :, z_index],
                                              target[:, :, z_index])[0]
            compute_num += 1
            Hausdorff_Distance += one_layer_HD
        return Hausdorff_Distance, round(Hausdorff_Distance / compute_num, 2)


def DSC(predict, target, threshold):
    epsilon = 1e-6
    assert (predict.shape == target.shape)
    if isinstance(predict, torch.Tensor) and isinstance(target, torch.Tensor):
        predict = torch.where(predict > threshold, 1, 0)
        TP = torch.sum(torch.multiply(predict, target)).item()
        FN = torch.sum(torch.multiply(1 - predict, target)).item()
        FP = torch.sum(torch.multiply(predict, 1 - target)).item()
        dsc = (2 * TP + epsilon) / (2 * TP + FN + FP + epsilon)
        if dsc > 1.0:
            print("A")
        return dsc
    elif isinstance(predict, np.ndarray) and isinstance(target, np.ndarray):
        predict = np.where(predict > threshold, 1, 0)
        TP = np.sum(np.multiply(predict, target))
        FN = np.sum(np.multiply(1 - predict, target))
        FP = np.sum(np.multiply(predict, 1 - target))
        dsc = (2 * TP + epsilon) / (2 * TP + FN + FP + epsilon)
        return dsc
    else:
        raise TypeError


def get_wrong_voxels(predict, target):
    if isinstance(predict, torch.Tensor) and isinstance(target, torch.Tensor):
        predict = predict.clone().detach().cpu().numpy()
        target = target.clone().detach().cpu().numpy()
    predict_correct = np.multiply(predict, target)
    predict_omit = np.multiply(1 - predict, target)
    predict_redundant = np.multiply(predict, 1 - target)
    return predict_correct, predict_omit, predict_redundant
